from datasets import load_dataset
from transformers import AutoTokenizer

# 加载分词器（必须与训练代码中的相同）
model_path = "./models/deepseek-ai/DeepSeek-R1-Distill-Qwen-1___5B"
tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenizer.pad_token = tokenizer.eos_token  # 确保与训练代码一致

# 加载原始数据
dataset = load_dataset("json", data_files={
    "train": "./data/train.jsonl",
    "validation": "./data/val.jsonl"
})


def process_fn(examples):
    # 步骤1：构建对话格式
    texts = [
        f"<|system|>\n你是一个助手</s>\n"
        f"<|user|>\n{prompt}</s>\n"
        f"<|assistant|>\n{response}</s>"
        for prompt, response in zip(examples['prompt'], examples['response'])
    ]

    # 步骤2：分词处理
    tokenized = tokenizer(
        texts,
        max_length=512,
        truncation=True,
        padding="max_length",
        return_tensors="pt"
    )

    # 步骤3：生成标签（因果语言建模专用处理）
    labels = tokenized["input_ids"].clone()
    # 将用户输入部分的标签设为-100（不计算损失）
    for i in range(len(texts)):
        user_start = texts[i].find("<|user|>")
        assistant_start = texts[i].find("<|assistant|>")
        labels[i, :assistant_start] = -100  # 仅保留助手的响应部分计算损失

    return {
        "input_ids": tokenized["input_ids"],
        "attention_mask": tokenized["attention_mask"],
        "labels": labels
    }


# 应用处理
dataset = dataset.map(
    process_fn,
    batched=True,  # 启用批量处理
    batch_size=1000,
    remove_columns=["prompt", "response"]  # 移除原始字段
)

# 保存处理后的数据集
dataset.save_to_disk("./data/processed_data")
